package com.blog.cmrpersonalblog.interceptor;

import cn.dev33.satoken.stp.StpUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.messaging.Message;
import org.springframework.messaging.MessageChannel;
import org.springframework.messaging.simp.stomp.StompCommand;
import org.springframework.messaging.simp.stomp.StompHeaderAccessor;
import org.springframework.messaging.support.ChannelInterceptor;
import org.springframework.messaging.support.MessageHeaderAccessor;
import org.springframework.stereotype.Component;

/**
 * WebSocket通道拦截器 - 用于消息级别的鉴权
 */
@Slf4j
@Component
public class WebSocketChannelInterceptor implements ChannelInterceptor {

    /**
     * 在消息发送前拦截
     */
    @Override
    public Message<?> preSend(Message<?> message, MessageChannel channel) {
        StompHeaderAccessor accessor = MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class);
        
        if (accessor != null) {
            StompCommand command = accessor.getCommand();
            
            // 处理CONNECT命令（连接时）
            if (StompCommand.CONNECT.equals(command)) {
                // 从STOMP头中获取token
                String token = accessor.getFirstNativeHeader("Authorization");
                
                if (token != null && token.startsWith("Bearer ")) {
                    token = token.substring(7);
                }
                
                // 如果STOMP头中没有token，尝试从会话属性中获取（握手时存入的）
                if (token == null || token.isEmpty()) {
                    Object sessionToken = accessor.getSessionAttributes().get("token");
                    if (sessionToken != null) {
                        token = sessionToken.toString();
                    }
                }
                
                try {
                    if (token != null && !token.isEmpty()) {
                        // 验证token
                        Object loginId = StpUtil.getLoginIdByToken(token);
                        if (loginId != null) {
                            Long userId = Long.parseLong(loginId.toString());
                            
                            // 将userId设置为STOMP用户
                            accessor.setUser(() -> userId.toString());
                            
                            log.info("WebSocket CONNECT成功: userId={}", userId);
                        } else {
                            log.warn("WebSocket CONNECT失败: token无效");
                            return null; // 拒绝连接
                        }
                    } else {
                        log.warn("WebSocket CONNECT失败: 缺少token");
                        return null; // 拒绝连接
                    }
                } catch (Exception e) {
                    log.error("WebSocket CONNECT鉴权失败", e);
                    return null; // 拒绝连接
                }
            }
            
            // 处理SUBSCRIBE命令（订阅时）
            if (StompCommand.SUBSCRIBE.equals(command)) {
                String destination = accessor.getDestination();
                String userId = accessor.getUser() != null ? accessor.getUser().getName() : null;

                if (userId == null) {
                    log.warn("WebSocket SUBSCRIBE失败: 用户未认证, destination={}", destination);
                    return null; // 拒绝订阅
                }

                // 验证用户只能订阅自己的通知
                if (destination != null && destination.startsWith("/user/")) {
                    // 支持两种订阅格式：
                    // 格式1: /user/queue/notifications（推荐，Spring自动路由）
                    // 格式2: /user/{userId}/queue/notifications（显式指定userId）

                    String[] parts = destination.split("/");

                    // 格式1: /user/queue/notifications (parts.length = 3)
                    // parts[0] = "", parts[1] = "user", parts[2] = "queue"
                    if (parts.length == 4 && "queue".equals(parts[2])) {
                        // 这是格式1，允许订阅（Spring会自动路由到当前用户）
                        log.debug("WebSocket SUBSCRIBE成功（格式1）: userId={}, destination={}", userId, destination);
                    }
                    // 格式2: /user/{userId}/queue/notifications (parts.length >= 4)
                    // parts[0] = "", parts[1] = "user", parts[2] = "{userId}", parts[3] = "queue"
                    else if (parts.length >= 4 && "queue".equals(parts[3])) {
                        String targetUserId = parts[2];

                        // 检查是否订阅自己的队列
                        if (!userId.equals(targetUserId)) {
                            log.warn("WebSocket SUBSCRIBE失败: 用户{}尝试订阅其他用户{}的通知, destination={}",
                                userId, targetUserId, destination);
                            return null; // 拒绝订阅
                        }
                        log.debug("WebSocket SUBSCRIBE成功（格式2）: userId={}, destination={}", userId, destination);
                    }
                    // 其他格式，允许订阅（可能是topic广播）
                    else {
                        log.debug("WebSocket SUBSCRIBE成功（其他格式）: userId={}, destination={}", userId, destination);
                    }
                } else {
                    // 非/user/开头的订阅（如/topic/），允许订阅
                    log.debug("WebSocket SUBSCRIBE成功（非用户队列）: userId={}, destination={}", userId, destination);
                }
            }
            
            // 处理DISCONNECT命令（断开连接时）
            if (StompCommand.DISCONNECT.equals(command)) {
                String userId = accessor.getUser() != null ? accessor.getUser().getName() : null;
                log.info("WebSocket DISCONNECT: userId={}", userId);
            }
        }
        
        return message;
    }

    /**
     * 消息发送后处理
     */
    @Override
    public void postSend(Message<?> message, MessageChannel channel, boolean sent) {
        // 可以在这里记录消息发送日志
    }

    /**
     * 消息发送完成后处理
     */
    @Override
    public void afterSendCompletion(Message<?> message, MessageChannel channel, boolean sent, Exception ex) {
        if (ex != null) {
            log.error("WebSocket消息发送失败", ex);
        }
    }
}

